import logging
import os
from datetime import datetime, timedelta
from typing import List, Any

import matplotlib.pyplot as plt
import numpy as np
import torch
from sklearn.metrics import mean_absolute_error, mean_squared_error

from config import config
from mydataloader import MyDataLoader
from mymodel import MyLSTMModel

# 获取配置
logger = logging.getLogger(__name__)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class PredictService:
    def __init__(self):
        self._dataloader = MyDataLoader()
        # 加载模型
        checkpoint = torch.load(os.path.join(config.MODELS_DIR, "geostyle_fit.pth"), map_location=device, weights_only=False)
        self._model = MyLSTMModel(
            num_cities=self._dataloader.get_cities_length(),
            num_attrs=self._dataloader.get_attributes_length()
        ).to(device)
        self._model.load_state_dict(checkpoint['model'])
        self._model.eval()

    def get_available_city_list(self) -> List[str]:
        """获取所有可用的城市列表"""
        return self._dataloader.get_available_city_list()

    def get_available_attributes(self, target_city_name: str) -> List[str]:
        """获取可用的属性列表"""
        return self._dataloader.get_available_attributes(target_city_name)

    def get_history_data(self, target_city_name: str, attribute: str) -> list[tuple[datetime, Any]]:
        """获取指定城市和属性的历史数据"""
        _, _, ts_data = self._dataloader.get_normed_history_data(target_city_name, attribute)
        # 反归一化
        ts_data = self._dataloader.denormalize_ts_data(target_city_name, attribute, ts_data)
        return ts_data

    def get_future_predict(self, target_city_name, attribute, steps):
        ts_data, city_id, attr_id, seqs = self._dataloader.get_predict_data(target_city_name, attribute)

        # 使用最后一个序列开始预测
        last_seq = seqs[-1]  # shape: [seq_len, feature_dim]
        last_date = ts_data[-1][0]
        delta_days = (ts_data[-1][0] - ts_data[-2][0]).days  # 计算时间步长

        future_predictions = []

        current_seq = last_seq.clone()  # 避免修改原数据
        city_tensor = torch.tensor([city_id], dtype=torch.long).to(device)
        attr_tensor = torch.tensor([attr_id], dtype=torch.long).to(device)
        for step in range(steps):
            input_tensor = current_seq.unsqueeze(0).to(device)

            with torch.no_grad():
                delta = self._model(input_tensor, city_tensor, attr_tensor).cpu().numpy()[0][0]

            last_value = current_seq[-1][0].item()
            predicted_value = last_value + delta

            # 更新日期
            last_date += timedelta(days=delta_days)
            future_predictions.append([last_date, float(predicted_value)])

            # 更新序列，去掉第一个，添加新预测值，保持shape一致
            new_step = current_seq[-1].clone()
            new_step[0] = float(predicted_value)  # 替换目标变量
            current_seq = torch.cat([current_seq[1:], new_step.unsqueeze(0)], dim=0)

        return future_predictions

    def predict_trend(self, target_city_name: str, attribute: str, steps: int) -> dict[str, list[tuple[datetime, Any]]]:
        """预测时尚趋势"""

        def generate_dates(base_date, step_days, n):
            """
            生成从基准时间开始按指定天数递增的日期列表

            参数:
                base_date (datetime): 基准时间
                step_days (int): 每次递增的天数（可为负数）
                n (int): 递增次数

            返回:
                List[datetime]: 生成的日期列表，长度为 n+1
            """
            return [base_date + timedelta(days=step_days * i) for i in range(n + 1)]

        def get_history_predict():
            ts_data, city_id, attr_id, seqs = self._dataloader.get_predict_data(target_city_name, attribute)

            historical_predictions = []

            city_tensor = torch.tensor([city_id], dtype=torch.long).to(device)
            attr_tensor = torch.tensor([attr_id], dtype=torch.long).to(device)

            for i, seq in enumerate(seqs):
                input_tensor = seq.unsqueeze(0).to(device)

                with torch.no_grad():
                    delta = self._model(input_tensor, city_tensor, attr_tensor).cpu().numpy()[0][0]

                last_value = seq[-1][0]
                predicted_value = last_value + delta
                historical_predictions.append([ts_data[i + self._dataloader.lookback][0], float(predicted_value)])

            return historical_predictions

        history = self.get_history_data(target_city_name, attribute)
        history_predict = self._dataloader.denormalize_ts_data(target_city_name, attribute, get_history_predict())
        future_predict = self._dataloader.denormalize_ts_data(target_city_name, attribute, self.get_future_predict(target_city_name, attribute, steps))

        return {
            "history": history,
            "all_time": generate_dates(history[0][0], abs((history[1][0] - history[0][0]).days), len(history) + steps - 1),
            "history_predict": history_predict,
            "future_predict": future_predict,
            "steps": steps
        }


def compute_history_metrics_with_plot(result, plot_title='History Prediction vs Actual'):
    """计算历史预测指标并绘图（含误差指标文本）"""
    # 构建字典
    history_dict = {t: v for t, v in result['history']}
    predict_dict = {t: v for t, v in result['history_predict']}

    # 找出时间戳交集并排序
    common_times = sorted(set(history_dict.keys()) & set(predict_dict.keys()))

    if not common_times:
        raise ValueError("No overlapping timestamps between history and history_predict.")

    y_true = np.array([history_dict[t] for t in common_times])
    y_pred = np.array([predict_dict[t] for t in common_times])

    # 指标计算
    mse = mean_squared_error(y_true, y_pred)
    rmse = np.sqrt(mse)
    mae = mean_absolute_error(y_true, y_pred)
    tra = np.mean(np.sign(np.diff(y_true)) == np.sign(np.diff(y_pred))) if len(y_true) >= 2 else np.nan

    # 绘图
    plt.figure(figsize=(10, 5))
    plt.plot(common_times, y_true, label='True', marker='o')
    plt.plot(common_times, y_pred, label='Predicted', marker='x')
    plt.title(plot_title)
    plt.xlabel('Time')
    plt.ylabel('Value')
    plt.legend()
    plt.xticks(rotation=45)
    plt.grid(True)

    # 添加指标文本
    textstr = f"MAE: {mae:.4f}\nMSE: {mse:.4f}\nRMSE: {rmse:.4f}\nTRA: {tra:.4f}"
    plt.gcf().text(0.75, 0.65, textstr, fontsize=10, bbox=dict(facecolor='white', alpha=0.6))

    plt.tight_layout()
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    plt.savefig(f"temp/output_plot_{timestamp}.png", dpi=300)
    plt.show()

    return {
        'MAE': mae,
        'MSE': mse,
        'RMSE': rmse,
        'TRA': tra
    }


if __name__ == "__main__":
    result = PredictService().predict_trend('Austin', 'GeoStyle___ageNA___genderNA___clothing_category__Dress', 12)
    print(compute_history_metrics_with_plot(result))
